#include "stdafx.h"
#include "MsgServer.h"
#include <windowsx.h>
#include <comdef.h>


#define WM_SERVICE_INVOKE	(WM_USER + 9999)
namespace msg{
	const wchar_t* wndclsName = L"EfetionAddinMessageWindow";

	wstring GetWndName(int id){
		wstring sid = boost::lexical_cast<wstring>(id);
		sid = wstring(L"EfetionAddinServer_") + sid;
		return sid;
	}
	typedef map<int, pair<HWND, OnService> > IdWndStore;
	
	IdWndStore* GetIdWnd(int flag){//flag: 0--get, 1--auto create, 2--auto delete
		static IdWndStore* all = 0;
		if(flag == 1){
			if(!all){
				all = new IdWndStore;
			}
		}else if(flag == 2){
			if(all && all->empty()){
				delete all;
				all = 0;
			}
		}
		return all;
	}

	bool WritVirMemString(HANDLE hProcess, LPVOID p, string s){
		DWORD sz = s.length();
		return WriteProcessMemory(hProcess, p, &sz, sizeof(sz), 0) &&
			WriteProcessMemory(hProcess, (char*)p + sizeof(sz), s.c_str(), sz, 0);
	}
	bool ReadVirMemString(HANDLE hProcess, LPVOID p, string& s){
		DWORD sz = 0;
		bool flag = false;
		if(ReadProcessMemory(hProcess, p, &sz, sizeof(sz), 0)){
			if(sz){
				s.resize(sz);
				if(ReadProcessMemory(hProcess, (char*)p + sizeof(sz), &s[0], sz, 0)){
					flag = true;
				}
			}else{
				s = "";
				flag = true;
			}
		}
		return flag;
	}

	LRESULT CALLBACK StaticWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam){
		if(uMsg == WM_SERVICE_INVOKE){
			IdWndStore* all = GetIdWnd(0);
			if(all){
				int id(wParam);
				IdWndStore::iterator it = all->find(id);
				if(it != all->end()){
					LPVOID p = (LPVOID)lParam;
					HANDLE hProcess = GetCurrentProcess();
					DWORD sz = 0;
					string request;
					if(ReadVirMemString(hProcess, p, request)){
						string reply = it->second.second(wParam, request);
						return WritVirMemString(hProcess, p, reply);
					}else{
						Log<<L"ReadProcessMemory Failed , size!"<<endl;
					}
				}
			}
			return FALSE;
		}
		return DefWindowProc(hWnd, uMsg, wParam, lParam);
	}

	HWND CreateMsgWindow(wstring name){
		static bool registed = false;
		if(!registed){
			WNDCLASSEXW wndcls = {0};
			wndcls.cbSize= sizeof(WNDCLASSEX);
			wndcls.style = CS_HREDRAW | CS_VREDRAW;
			wndcls.lpfnWndProc = &StaticWndProc;
			wndcls.cbClsExtra = 0;
			wndcls.cbWndExtra = 0;
			wndcls.hInstance = GetModuleHandle(0);
			wndcls.hbrBackground = (HBRUSH)(COLOR_WINDOW +1);
			wndcls.hCursor = LoadCursor(NULL, IDC_ARROW);
			wndcls.hIcon = 0;
			wndcls.hIconSm = 0;
			wndcls.lpszClassName = wndclsName;
			wndcls.lpszMenuName =0;
			registed= RegisterClassEx(&wndcls) != 0;
		}
		if(!registed){
			return 0;
		}
		return ::CreateWindowEx(0, wndclsName, name.c_str(), WS_POPUP, 0, 0, 0, 0, 0, 0, GetModuleHandle(0), 0);
	}

	bool Open(int id, OnService service){
		wstring name = GetWndName(id);
		HWND w = ::FindWindowEx(0, 0,wndclsName, name.c_str());
		if(w){
			Log<<L"Try Create Window With Exist Name!!"<<name<<endl;
			return false;
		}
		HWND wnd = CreateMsgWindow(name);
		if(wnd){
			GetIdWnd(1)->insert(make_pair(id, make_pair(wnd, service) ));
			Log<<"Open "<<id<<" Success."<<endl;
			return true;
		}
		//else
		Log<<L"Create Message Window Failed!"<<endl;
		//Log<<Util::GetLastErrorDesc()<<endl;
		return false;
	}
	void Close(int id){
		wstring name = GetWndName(id);
		IdWndStore* all = GetIdWnd(0);
		if(all){
			HWND w = ::FindWindowEx(0, 0,wndclsName, name.c_str());
			if(w){
				IdWndStore::iterator it = all->find(id);
				if(it != all->end()){
					HWND wnd = it->second.first;
					if(w == wnd){
						DestroyWindow(w);
						all->erase(id);
						Log<<"Close Success! Name: "<<name<<endl;
					}
				}else{
					Log<<"Try Close Window In Other Process! "<<name<<endl;
				}
			}else{
				all->erase(id);
				Log<<L"NO Window Found In All Process! Name : "<<name<<endl;
			}
			
		}else{
			Log<<L"NO Window Found In Current Process! Name :"<<name<<endl;
		}
		GetIdWnd(2);
	}
	string Invoke(int id, string text){
		wstring name = GetWndName(id);
		string reply;
		HWND w = ::FindWindowEx(0, 0,wndclsName, name.c_str());
		if(w){
			DWORD pid = 0;
			GetWindowThreadProcessId(w, &pid);
			HANDLE hProcess = OpenProcess(PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_VM_WRITE, FALSE, pid);
			if(hProcess && hProcess != INVALID_HANDLE_VALUE){
				DWORD sz = text.length();
				LPVOID p = VirtualAllocEx(hProcess, 0, sz + sizeof(sz), MEM_COMMIT, PAGE_READWRITE);
				if(p){
					if( WritVirMemString(hProcess, p, text)){
						LRESULT lr = SendMessage(w, WM_SERVICE_INVOKE, id, LPARAM(p));
						if(lr){
							
							if(!ReadVirMemString(hProcess, p, reply)){
								reply = "";
								Log<<L"Invoke.ReadVirMemString Failed: "<<name<<endl;
							}
						}else{
							Log<<L"SendMessage Return 0!"<<name<<endl;
						}
					}else{
						Log<<L"WriteProcessMemory Failed: "<<name<<endl;
					}
					VirtualFreeEx(hProcess, p, 0, MEM_RELEASE);
				}else{
					Log<<L"VirtualAllocEx Failed: "<<name<<endl;
				}
			}else{
				Log<<L"Open Process Failed!"<<name<<endl;
			}
		}else{
			Log<<L"No Window Named: "<<name<<endl;
		}
		return reply;
	}

}